from ImageCorrector import ImageCorrector
import cv2
import numpy as np

class DisparityGenerator:
    def __init__(self,imageCorrector):
        self.imageCorrector = imageCorrector#调用ImageCorrector类
        #左右校正之后的图转为灰度图
        self.imgL = cv2.cvtColor(self.imageCorrector.left_rectified,cv2.COLOR_BGR2GRAY)
        self.imgR = cv2.cvtColor(self.imageCorrector.right_rectified,cv2.COLOR_BGR2GRAY)

    def generate(self,num,blockSize, show=False):
        # 创建窗口
        cv2.namedWindow('SGBM_processed')

        # 左右立体匹配参数 左视差图
        stereo = cv2.StereoSGBM_create(
            minDisparity=0,
            numDisparities=16 * num,  # numDisparities 必须是16的倍数
            blockSize=blockSize,  # blockSize 必须是奇数
            P1=8 * blockSize ** 2,
            P2=32 * blockSize ** 2,
            disp12MaxDiff=1,
            uniquenessRatio=10,
            speckleWindowSize=100,
            speckleRange=32
        )

        # 计算视差图
        disparity_img = stereo.compute(self.imgL, self.imgR).astype(np.float32) / 16.0  # 左视差图
        # 归一化视差图，便于显示
        disparity_normalized = cv2.normalize(disparity_img,None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX)
        disparity_normalized = disparity_normalized.astype(np.uint8)

        # 转为伪彩色显示
        disparity_color = cv2.applyColorMap(disparity_normalized, cv2.COLORMAP_JET)

        """
        show == True 代表我需要显示，并且持久化深度图和颜色图 调试代码使用
        show == False 代表不需要显示，一般正常调用选择此选项 也是默认的选项
        """
        if show == True:
            # 显示原始视差图
            cv2.imshow('SGBM_processed', disparity_color)
            # 保存生成的深度图和左半边图
            cv2.imwrite('depth_map.png')  # 保存深度图（视差图）
            cv2.imwrite('original.png', self.imageCorrector.img[0:480, 0:640])  # 保存原始的左半边图
            print("视差图和左半边图已保存！")
            # 按 ESC 键退出
            cv2.waitKey(0)
            cv2.destroyAllWindows()
        else:
            return ( disparity_img,self.imageCorrector.img[0:480, 0:640])
    """
    不同照片可找到合适的参数，一般不用
    """
    def parameterSelection(self):

        cv2.namedWindow('SGBM processing')
        cv2.createTrackbar('num', 'SGBM processing', 2, 10, lambda x: None)  # 滑块：numDisparities 参数
        cv2.createTrackbar('blockSize', 'SGBM processing', 5, 255, lambda x: None)  # 滑块：blockSize 参数
        while True:
            # 获取 SGBM 参数
            num = cv2.getTrackbarPos('num', 'SGBM processing')
            blockSize = cv2.getTrackbarPos('blockSize', 'SGBM processing')

            # 确保 blockSize 为奇数且 >= 5
            if blockSize % 2 == 0:
                blockSize += 1
            if blockSize < 5:
                blockSize = 5

            # 左右立体匹配参数
            stereo = cv2.StereoSGBM_create(
                minDisparity=0,
                numDisparities=16 * num,  # numDisparities 必须是16的倍数
                blockSize=blockSize,  # blockSize 必须是奇数
                P1=8 * blockSize ** 2,
                P2=32 * blockSize ** 2,
                disp12MaxDiff=1,
                uniquenessRatio=10,
                speckleWindowSize=100,
                speckleRange=32
            )

            # 计算视差图
            disparity_img = stereo.compute(self.imgL, self.imgR).astype(np.float32) / 16.0  # 左视差图
            # 归一化视差图，便于显示
            disparity_normalized = cv2.normalize( disparity_img,None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX)
            disparity_normalized = disparity_normalized.astype(np.uint8)

            # 转为伪彩色显示
            disparity_color = cv2.applyColorMap(disparity_normalized, cv2.COLORMAP_JET)

            cv2.imshow('SGBM', disparity_color)

            # 按 ESC 键退出并保存参数
            if cv2.waitKey(1) & 0xFF == 27:
                with open("parameters.txt", "w") as file:
                    file.write(f"num={num}\n")
                    file.write(f"blockSize={blockSize}\n")
                print(f"最终选择的参数已保存到 parameters.txt: num={num}, blockSize={blockSize}")
                break
        cv2.destroyAllWindows()